import pickle
import yaml
import os
import csv

import numpy as np
import matplotlib.pyplot as plt

from tqdm import tqdm
from tabulate import tabulate
from loguru import logger

# ANSI转义码定义颜色
RESET = "\033[0m"  # 重置/默认颜色
BOLD = "\033[1m"   # 粗体
RED = "\033[31m\033[1m"   # 红色粗体
GREEN = "\033[32m\033[1m" # 绿色粗体
BLUE = "\033[34m\033[1m"  # 蓝色粗体

# 为关键词设置颜色的函数
def colorize(text, color_code):
    return f"{color_code}{text}{RESET}"


def print_help(filename, COMMANDS):
    usage_text = colorize("Usage", GREEN) + ": python "+ colorize(filename,BOLD) +" " + colorize("<COMMAND>", BLUE)
    version_text = colorize("Version", GREEN) + ": 1.0"
    commands_header = colorize("Available commands:", GREEN)
    help_option = colorize("-h, --help", BLUE) + "       Show this help message"
    
    print(usage_text)
    print(version_text)
    print("\n" + commands_header)
    for cmd, desc in COMMANDS.items():
        cmd_text = colorize(cmd, BLUE)  # 将命令设置为蓝色加粗
        print(f"  {cmd_text.ljust(25)} {desc}")
    print("\nOptions:")
    print("  " + help_option)


def save_dataset(data_item, filename):
    """
    将数据集保存到指定的文件中。

    参数:
    data_item (any): 要保存的数据集。
    filename (str): 保存数据集的文件名。
    """
    print(f"Waiting for saving dataset to {filename}...")
    try:
        with open(filename, 'wb') as file:
            pickle.dump(data_item, file, protocol=pickle.HIGHEST_PROTOCOL)
        print(f"Dataset saved to {filename}")
    except IOError as e:
        print(f"Failed to save dataset: {e}")

def save_dataset_chunk(data_item, filename):
    """
    将数据集保存到指定的文件中。

    参数:
    data_item (any): 要保存的数据集。
    filename (str): 保存数据集的文件名。
    """
    tqdm.write(f"\nWaiting for saving dataset to {filename}...")
    with open(filename, 'wb') as file:
        while True:
            data_item_empty = []
            data_item = yield data_item_empty
            pickle.dump(data_item, file, protocol=pickle.HIGHEST_PROTOCOL)
            # tqdm.write(f"Dataset saved to {filename}")
    

def load_dataset(filename):
    """
    从指定的文件中加载数据集。

    参数:
    filename (str): 要加载数据集的文件名。

    返回:
    any: 加载的数据集。
    """
    print(f"Waiting for loading dataset from {filename}...")
    try:
        with open(filename, 'rb') as file:
            data_set = pickle.load(file)
        print(f"Dataset loaded from {filename}")
        return data_set
    except (IOError, pickle.PickleError) as e:
        print(f"Failed to load dataset: {e}")
        return None
    
def load_dataset_chunk(filename):
    """
    从指定的文件中加载数据集。需要与save_dataset_chunk配合使用。

    参数:
    filename (str): 要加载数据集的文件名。

    返回:
    any: 加载的数据集。
    """
    tqdm.write(f"\nWaiting for loading dataset from {filename}...")
    with open(filename, 'rb') as file:
        while True:
            try:
                data_set = pickle.load(file)
                tqdm.write(f"Dataset loaded from {filename}")
                yield data_set
            except EOFError:
                break

def count_chunks(filename):
    count = 0
    with open(filename, 'rb') as f:
        while True:
            try:
                pickle.load(f)  # 尝试加载一个块
                count += 1
            except EOFError:    # 捕获文件结束异常
                break
    return count

def save_loss_curve(loss_values, output_path='loss_curve.png'):
    plt.figure(figsize=(10, 5))
    plt.plot(loss_values, 'b-', label='Training Loss')
    plt.title('Training Loss Curve')
    plt.yscale('log')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.grid(True)
    plt.legend()
    plt.savefig(output_path)  # 保存为 PNG
    # plt.show()
    plt.close()

def save_last_epoch_curve(X_list, y_list, pred_list, last_epoch_loss_value, output_path='last_epoch_curve.png'):
  
    X_np = X_list[0].detach().numpy()
    y_np = y_list[0].detach().numpy()
    pred_np = pred_list[0].detach().numpy()


    omega_x_diff2 = (y_np[:,0] - pred_np[:,0])**2
    omega_y_diff2 = (y_np[:,1] - pred_np[:,1])**2
    loss_cal_manual = 0
    for i in range(len(omega_x_diff2)):
        loss_cal_manual += (omega_x_diff2[i] + omega_y_diff2[i])/(len(omega_x_diff2) + len(omega_y_diff2))
        

    with open('last_epoch_data.csv', mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['X', 'y', 'pred','loss','loss_cal_manual'])
        # for i in range(len(X_list)):
        writer.writerow([X_list[0].detach().numpy(), 
                         y_list[0].detach().numpy(), 
                         pred_list[0].detach().numpy(), 
                         last_epoch_loss_value[0],
                         loss_cal_manual])
    
def print_table(tabel_data, table_title):
    # tablefmt options: simple_grid rounded_grid heavy_grid mixed_grid
    print(f"{table_title}:")
    print(tabulate(tabel_data, headers=["Item", "Length"], tablefmt="rounded_grid"))

def load_yaml(yaml_name):
    """Load YAML configuration file.
    
    Args:
        yaml_name: Name of the YAML file (without extension)
        
    Returns:
        dict: Parsed YAML data or empty dict if file not found/error
    """
    yaml_path = f'config_yaml/{yaml_name}.yaml'
    
    if not os.path.exists(yaml_path):
        logger.error(f"YAML file not found: {yaml_path}")
        return None
    
    try:
        with open(yaml_path, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
            return data if data is not None else {}
    except (IOError, yaml.YAMLError) as e:
        logger.error(f"Error loading YAML file {yaml_path}: {str(e)}")
        return None